import json
import sys
import argparse
import fcntl
import subprocess
import os
import time


def get_host_ip():
    if sys.version_info >= (3, 5):
        result = subprocess.run(['hostname', '-i'], stdout=subprocess.PIPE)
        return result.stdout.decode('utf-8').strip()
    else:
        process = subprocess.Popen(['hostname', '-i'], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        stdout, stderr = process.communicate()
    return stdout.decode('utf-8').strip()


def get_device_ips():
    device_ips = []
    for i in range(8):
        if sys.version_info >= (3, 5):
            result = subprocess.run(['hccn_tool', '-i', str(i), '-ip', '-g'], stdout=subprocess.PIPE)
            output = result.stdout.decode('utf-8').strip()
        else:
            process = subprocess.Popen(['hccn_tool', '-i', str(i), '-ip', '-g'], stdout=subprocess.PIPE,
                                       stderr=subprocess.PIPE)
            stdout, stderr = process.communicate()
            output = stdout.decode('utf-8').strip()
        ipaddr = output.split('\n')[0].split(':')[1].strip()
        device_ips.append(ipaddr)
    return device_ips


def is_previous_server_complete(data, server_index):
    if server_index == 1:
        return True

    previous_server_end_rank = (server_index - 1) * 8 - 1
    rank_ids = set()

    for server in data.get("server_list", []):
        for device in server.get("device", []):
            rank_ids.add(int(device["rank_id"]))

    for rank in range(previous_server_end_rank - 7, previous_server_end_rank + 1):
        if rank not in rank_ids:
            return False
    return True


def wait_for_previous_server(json_file, server_index, check_interval=5):
    while True:
        if not os.path.exists(json_file):
            print("File %s does not exist. Waiting..." % json_file)
            time.sleep(check_interval)
            continue

        with open(json_file, 'r') as f:
            try:
                data = json.load(f)
                if is_previous_server_complete(data, server_index):
                    print("Previous server completed. Proceeding...")
                    return
                else:
                    print("Previous server not completed. Waiting...")
            except json.JSONDecodeError:
                print("File %s is not valid JSON. Waiting..." % json_file)
        time.sleep(check_interval)


def update_json(json_file, server_id, container_ip, server_count, server_index):
    if server_index != 1:
        wait_for_previous_server(json_file, server_index)
    if os.path.exists(json_file):
        with open(json_file, 'r') as f:
            fcntl.flock(f, fcntl.LOCK_EX)
            data = json.load(f)
    else:
        data = {
            "version": "1.0",
            "server_count": str(server_count),
            "server_list": [],
            "status": "completed"
        }

    if data.get("server_count") == str(server_count):
        pass
    else:
        data["server_count"] = str(server_count)

    rank_id_start = (server_index - 1) * 8
    server_found = False
    devices = [{"device_id": str(i), "device_ip": ip, "rank_id": str(rank_id_start + i)} for i, ip in
               enumerate(get_device_ips())]

    for server in data["server_list"]:
        if server["server_id"] == server_id:
            server_found = True
            if container_ip:
                server["container_ip"] = container_ip
            server["device"] = devices
            break

    if not server_found:
        new_server = {
            "server_id": server_id,
            "container_ip": container_ip if container_ip else server_id,
            "device": devices,
            "host_nic_ip": "reserve"
        }
        data["server_list"].append(new_server)

    data["server_list"] = sorted(data["server_list"], key=lambda x: int(x["device"][0]["rank_id"]))

    with open(json_file, 'w') as f:
        fcntl.flock(f, fcntl.LOCK_EX)
        json.dump(data, f, indent=4)
        fcntl.flock(f, fcntl.LOCK_UN)
        if server_index == server_count:
            os.chmod(json_file, 0o640)


def main():
    parser = argparse.ArgumentParser(description="Update JSON file with server information.")
    parser.add_argument('--server_id', type=str, required=True)
    parser.add_argument('--container_ip', type=str)
    parser.add_argument('--server_count', type=int, default=4)
    parser.add_argument('--server_index', type=int, required=True)
    parser.add_argument('--json_file', type=str, default='new_rank_table_file.json')

    args = parser.parse_args()

    server_id = args.server_id
    container_ip = args.container_ip if args.container_ip else server_id

    update_json(args.json_file, server_id, container_ip, args.server_count, args.server_index)


if __name__ == "__main__":
    main()
